# 为了做ppt，不保存目标框，只保存拟合的椭圆。面积换算成实际值
# YOLOv3 🚀 by Ultralytics, GPL-3.0 license
"""
Run inference on images, videos, directories, streams, etc.

Usage:
    python detect_ltl.py --save-txt --save-conf --conf-thres 0.5 --name 'data0716_exp6_thin_overlopdetect' --line-thickness 1 --hide-labels --weights './runs/train/exp6/weights/best.pt' --img 512 --source './data/data0716/img_716_thin'
    python detect.py --save-txt --save-conf --save-crop --conf-thres 0.5 --name 'databubble0616-crop' --line-thickness 1 --hide-labels --weights 'runs/train/databubble0604_3/weights/best.pt' --img 512 --source 'D:\\whitebubble\\machinelearning\\dataset\\data0616\\img_616'
"""

import argparse
import os
import sys
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from utils.ellipsefit_ltl import ellipse_detect

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.common import DetectMultiBackend
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
                           increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, time_sync


@torch.no_grad()
def run(weights=ROOT / 'runs/train/exp6/weights/best.pt',  # model.pt path(s)
        source=ROOT/'data/data0716/writing',  # file/dir/URL/glob, 0 for webcam
        img=512,  # inference size (pixels)
        conf_thres=0.25,  # confidence threshold
        iou_thres=0.45,  # NMS IOU threshold
        max_det=1000,  # maximum detections per image
        device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
        save_txt=True,  # save results to *.txt
        save_conf=True,  # save confidences in --save-txt labels
        save_crop=True,  # save cropped prediction boxes

        classes=None,  # filter by class: --class 0, or --class 0 2 3
        agnostic_nms=False,  # class-agnostic NMS

        update=False,  # update all models
        project=ROOT / 'runs/detect',  # save results to project/name
        name='data0716_exp6_writing',  # save results to project/name
        exist_ok=False,  # existing project/name ok, do not increment
        line_thickness=1,  # bounding box thickness (pixels)
        hide_labels=True,  # hide labels
        hide_conf=True,  # hide confidences

        dnn=False,  # use OpenCV DNN for ONNX inference
        ):
    imgsz = [img, img]  # expand
    source = str(source)
    save_img = True  # save inference images

    # Directories
    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
    (save_dir / 'labels').mkdir(parents=True, exist_ok=True)  # make dir

    # Load model
    device = select_device(device)
    model = DetectMultiBackend(weights, device=device, dnn=dnn)
    stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Dataloader
    dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
    bs = 1  # batch_size

    # Run inference
    if pt and device.type != 'cpu':
        model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters())))  # warmup
    dt, seen = [0.0, 0.0, 0.0], 0
    for path, im, im0s, vid_cap, s in dataset:
        t1 = time_sync()
        im = torch.from_numpy(im).to(device)
        im = im.float()
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
        t2 = time_sync()
        dt[0] += t2 - t1

        # Inferencel's
        pred = model(im)
        t3 = time_sync()
        dt[1] += t3 - t2

        # NMS 非极大值抑制
        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
        dt[2] += time_sync() - t3

        # Process predictions
        for i, det in enumerate(pred):  # per image
            area_total = 0
            seen += 1
            p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # im.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
            s += '%gx%g ' % im.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            imc = im0.copy() if save_crop else im0  # for save_crop
            img_save = im0.copy()  # 最后保存要保存img_save，避免在原图上修改时，改变crop，使得检测椭圆时图片中有拟合的其他椭圆/像素值
            annotator = Annotator(img_save, line_width=line_thickness, example=str(names))
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                # for *xyxy, conf, cls in reversed(det):
                for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop:  # Add bbox to image
                        c = int(cls)  # integer class
                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                        # annotator.box_label(xyxy, label, color=colors(c, True))
                        x1 = (xywh[0]-xywh[2]/2)*img
                        x2 = (xywh[0]+xywh[2]/2)*img
                        y1 = (xywh[1]-xywh[3]/2)*img
                        y2 = (xywh[1]+xywh[2]/2)*img
                        crop = imc[int(y1):int(y2), int(x1):int(x2)]
                        if abs(x1-x2)<15 or abs(y1-y2)<15: # 动态的上采样倍数吧
                            times = 3
                        elif abs(x1-x2)<30 or abs(y1-y2)<30:
                            times = 2
                        else:
                            times = 1
                        _ellipse, area = ellipse_detect(crop, times)
                        cv2.ellipse(img_save, [[_ellipse[0][0]/times+x1,_ellipse[0][1]/times+y1],[_ellipse[1][0]/times,_ellipse[1][1]/times],_ellipse[2]], (100, 0, 0), 1) # 椭圆坐标对应需要改 xy是反的
                        area_total += area
                        # cv2.putText(img_save, '%d pixels'%area, (int(x1)-5, int(y1)-5), cv2.FONT_HERSHEY_PLAIN, 0.65,(255, 255, 255))
            area_total = area_total*0.060606
            cv2.putText(img_save, '%d mm2' % area_total, (10, img-20), cv2.FONT_HERSHEY_COMPLEX, 1.5, (255, 255, 255))
            cv2.putText(img_save, '%d bubbles' % det.shape[0], (10, img - 60), cv2.FONT_HERSHEY_COMPLEX, 1.5,(255, 255, 255))
            # Print time (inference-only)
            LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')

            # Save results (image with detections)
            if save_img:
                cv2.imwrite(save_path, img_save)


    # Print results
    t = tuple(x / seen * 1E3 for x in dt)  # speeds per image
    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
    if save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
    if update:
        strip_optimizer(weights)  # update model (to fix SourceChangeWarning)

run()

